from class_family import *

class Equation(object):
	'''
	an Equation object to check equivalence of left- and righ-hand-sides
	'''
	def __init__(self, left, right, weight, disp = False):

		self.left = left
		self.right = right
		self.weight = weight
		#self.weight.shape = (-1,1)
		self.error = self.left - np.dot(self.right, self.weight)

		# not for matrix weights
		self.right_term = self.right * self.weight.ravel()

		self.termsNoWeight = np.hstack( (self.left , self.right ) )

		self.termsWeight = np.hstack( (self.left , self.right_term) )

		self.checkError(disp = disp)
		

	def checkError(self , disp = False):

		e = np.mean( np.abs(self.error) )
		base = np.sum( np.abs(self.left) )

		if disp:
			print('-'*50)
			print( 'Check Equation Error:' )
			if e < ErrorTol * base:
				print('Decomposition Succeeded! Error: {:.4f}'.format(e))
			else:
				raise ValueError('Decomposition Failed! Error: {:.4f}'.format(e))
			print('-'*50)

	def getRow(self):

		temp = np.hstack( (self.left , self.right ) )

		self.row = np.hstack( ( temp.ravel(), self.weight.ravel() ) )


class EquationPi(Equation):

	def __init__(self, left, right, disp):
		self.left = left
		self.right = right
		self.error = self.left - (self.right[0] - self.right[1])
		self.weight = []
		self.checkError(disp = disp)


class Decomposition(object):
	'''
	decompose the QQ effect to the price, substitution, and rationing income effects
	start from a restricted q2 (r2) model, at a given value of n
	'''
	def __init__(self, FR2, disp=False):

		assert type(FR2) is type(FamilyRation()), 'FR2 should be FamilyRation object!'

		self.FR2 = FR2


	def nxset(self, n = [1], x0 = [0.1,0.1], disp = False):
		'''
		FR2: restricted q2 model, our cornerstone
		
		'''
		self.FR2.n = n
		self.FR2.x = x0
			
			# update the derivatives for given values of q, s, n
		self.FR2.bindParameters()
		self.FR2.updateNumerics()

			# optimize by choosing x
		if disp:
			print('Initialization of the restricted q2 model')

		self.FR2.dualOptimizeOffer(umpfirst = True, disp = disp)

	def update(self, disp = False):
		'''
		FR2: restricted q2 model, our cornerstone
		'''
		if disp:
			print('Initialization of the restricted q2 model')

		self.FR2.dualOptimizeOffer(umpfirst = True, disp = disp)

	def run(self, disp = False):

		self.FR2umpemp(disp = disp)
		self.FR2toFR1(disp = disp)
		self.FR1toFU1(disp = disp)
		#self.FR2toFU2(disp = disp)

		self.getData()

	def FR2umpemp(self, disp = False):
		'''
		restricted q2 model, ump and emp
		'''
		if disp:
			print('UMP and EMP of the r2 model')

		left = self.FR2.dxdn
		right = np.hstack((
			self.FR2.dxdn_emp,
			self.FR2.yeffect,
			))
		weight = np.array([1, self.FR2.spe_pe[0]])
		weight.shape = (-1,1)

		self.R2umpemp = Equation(
			left = left,
			right = right,
			weight = weight,
			disp = disp
			)

	def FR2toFR1(self, disp = False):
		'''
		r2 and r1 model
		'''
		if disp:
			print('Decomposition: r2 -> r1')

		# initializing the restricted q1 model
		self.FR1 = self.FR2.getR1()

		# get the INDUCED restricted q1 model
		self.FR1.piset(self.FR2.pi_r1)
		self.FR1.tuset(self.FR2.u)
		self.FR1.xset(self.FR2.x) # a wise initial value~haha~

		self.FR1.dualOptimizeOffer(umpfirst = False)


		if disp:
			print('-'*50)
			print( 'Check Induced Model: restricted q2 and q1' )
			error = np.mean( self.FR2.yeffect - self.FR1.yeffect )
			if abs(error) < ErrorTol:
				print('Test 1 Succeeded! Y effects are the same!')
			else:
				raise ValueError('Induced Change Failed! Y effects are not the same')

			error = np.mean( np.array(self.FR1.x) - np.array(self.FR2.x) )
			if abs(error) < ErrorTol:
				print('Test 2 Succeeded! X are the same!')
			else:
				raise ValueError('Induced Change Failed! X are not the same')
			print('-'*50)

		# decomposition

		dxdpi_q = self.FR1.slutsky_emp[:, :1]

		left = self.FR2.dxdn_emp
		right =  np.hstack((
			self.FR1.dxdn_emp,
			dxdpi_q
			))
		weight = np.array([1, self.FR2.pi[-1]])
		weight.shape = (-1,1)

		self.R2R1eq = Equation(left = left,
			right = right,
			weight = weight,
			disp = disp
			)

	def RtoU(self, FR, FU, disp = False):
		'''
		Neary & Roberts 1980 decomposition
		'''

		# Step 1: get the induced unrestricted model

		assert FR.dualOpt, 'Restricted Family Not Optimized Yet!'
		if disp:
			print('Get the induced unrestricted, NR 1980')
		FU.piset(FR.spi)
		FU.tuset(FR.u) # use u to induce, not tu!!!!!!!!!!
		#I have spendt a huge amount of time to find this error!!!

		FU.xset(FR.n + FR.x) # get a wise initial value? haha~
		FU.dualOptimize(umpfirst = False)

		if disp:
			print('-'*50)
			print( 'Check Induced Model: restricted to unrestricted' )
			error = np.mean( np.array(FR.n + FR.x) - np.array(FU.x) )
			if abs(error) < ErrorTol:
				print('Test 1 Succeeded! Commodity Bundles X are the same!')
			else:
				print('Induced Change Failed! X are not the same.')
				print(FR.n + FR.x)
				print(FU.x)
				raise
			print('-'*50)

			print('-'*50)
			print( 'Check Induced Model: restricted to unrestricted' )
			error = FR.sy - FU.y
			if abs(error) < ErrorTol:
				print('Test 2 Succeeded! Income level sy has been induced!')
			else:
				print('Induced Change Failed! Income level sy has not been induced.')
				print(FR.sy)
				print(FU.y)
				raise
			print('-'*50)

		# Step 2: change of n
		if disp:
			print('Decomposition, change of n')
		FU.dxedpi_n = FU.csemp[ :-1, :1 ]
		FU._dndpi_n = FU.dxedpi_n[ :1 ]
		FU._dxdpi_n = FU.dxedpi_n[ 1: ]

		left = FR.dxdn_emp
		right = FU._dxdpi_n
		weight = np.linalg.inv(FU._dndpi_n)

		RUeqn = Equation(
			left = left,
			right = right,
			weight = weight,
			disp = disp
			)

		# step 3: change of y
		if disp:
			print('Decomposition, change of y')
		FU._dndy = FU.yeffect[:1, :]
		FU._dxdy = FU.yeffect[1:, :]

		left = FR.yeffect
		right = np.hstack((
			FU._dxdy,
			np.dot( FR.dxdn_emp, FU._dndy ),
			))
		weight = np.array([1,-1])
		weight.shape = (-1,1)

		RUeqy = Equation(
			left = left,
			right = right,
			weight = weight,
			disp = disp
			)

		# step 4: change of pi
		if disp:
			print('Decomposition, change of pi')
		FR.dxdpif_emp = FR.slutsky
		FU._dxdpi_x_emp = FU.slutsky[1:,1:]
		FU._dxdpi_n_emp = FU.slutsky[1:, :1]
		FU._dndpi_n_emp = FU.slutsky[:1, :1]
		FU._dndpi_x_emp  = FU.slutsky[:1, 1:]


		temp = np.dot(FU._dxdpi_n_emp, np.linalg.inv(FU._dndpi_n_emp) )
		temp = np.dot(temp, FU._dndpi_x_emp)


		left = FR.dxdpif_emp
		right = np.array((
			FU._dxdpi_x_emp,
			temp
			))

		RUeqpi = EquationPi(
			left = left,
			right = right,
			disp = disp
			)

		return RUeqn, RUeqy, RUeqpi

	def FR2toFU2(self, disp = False):
		'''
		restricted q2 model to unrestricted q2 model
		'''
		if disp:
			print('Decomposition: r2 -> u2 model')

		self.FU2 = self.FR2.getU2()

		self.R2U2eqn, self.R2U2eqy, self.R2U2eqpi = \
			self.RtoU(FR = self.FR2, FU = self.FU2, disp = disp)

	def FR1toFU1(self, disp = False):
		'''
		restricted q1 model to unrestricted q1 model
		'''
		if disp:
			print('Decomposition: r1 -> u1 model')

		self.FU1 = self.FR2.getU1()

		self.R1U1eqn, self.R1U1eqy, self.R1U1eqpi = \
			self.RtoU(FR = self.FR1, FU = self.FU1, disp=disp)

	def checkDecomp(self, disp = False):

		if disp:
			print('Decomposition')

		self.FU1._dndy = self.FU1.yeffect[:1, :]
		self.FU1._dxdy = self.FU1.yeffect[1:, :]

		left = self.FR2.dxdn
		right =  np.hstack((
			self.FR1.slutsky_emp[:, :1], # price effect
			self.FR1.dxdn_emp, # substitution effect
			self.FU1._dxdy # can also use self.FR2.yeffect
			))

		weight1 = self.FR2.pi[-1]
		weight2 = 1 - self.FR2.spe_pe[0] * self.FU1._dndy[0,0]
		weight3 = self.FR2.spe_pe[0]
		weight = np.array( [ weight1, weight2, weight3] )
		weight.shape = (-1,1)
		EquationDecomp = Equation(
			left = left,
			right = right,
			weight = weight,
			disp = disp
			)

		print('Left:')
		print(EquationDecomp.left)

		print('Right:')
		print(EquationDecomp.right_term)


	def getData(self):
		'''
		compute important derivatives and prices
		store in xData, pData, dnData
		'''

		lenx = self.FR2.lenx
		strx = self.FR2.strx

		self.strx = strx

		# commodity bundle

		xData = np.array([self.FR2.x])

		col_label = strx

		self.xData = pd.DataFrame(
			data = xData,
			columns = col_label
			)

		# dx: derivatives of x (q and s)
		temp = np.hstack((
			self.FR2.dxdn,
			self.FR1.dxdn_emp,
			self.FU1._dxdpi_x_emp[:, :1],
			self.FR1.slutsky_emp[:, :1],
			self.FR1.yeffect,
			self.FU1._dxdy
			))

		col_label = ['r2dxdn', 'r1dxdn_emp', 'u1_dxdpi_q_emp',
		 'r1dxdp_q_emp', 'r1dxdy', 'u1_dxdy']

		self.dxData	= {}
		for i in range(lenx):
			sx = strx[i]
			data = temp[i:i+1, :]
			self.dxData[sx] = pd.DataFrame(
				data = data,
				columns = col_label
				)

		# dn: derivatives of n
		dndy = self.FU1._dndy

		dnData = np.hstack((
			self.FU1._dndy,
			self.FU1._dndpi_n_emp,
			self.FU1._dndpi_x_emp
			))
		col_label = ['u1_dndy', 'u1_dndpi_n_emp',
			'u1_dndpi_q_emp', 'u1_dndpi_s_emp'
			]

		self.dnData = pd.DataFrame(
			data = dnData,
			columns = col_label
			)

		# shadow prices
		sp_n_p_n = self.FR2.spe_pe[0]
		spi_n = self.FR2.spi[0]
		sp_n  = self.FR2.spe[0]
		pi_nq = self.FR2.pi[-1]

		pData = np.array( [[sp_n_p_n, spi_n, sp_n, pi_nq ]] )
		col_label = ['sp_n_p_n', 'spi_n', 'sp_n', 'pi_nq' ]

		self.pData = pd.DataFrame(
			data = pData,
			columns = col_label
			)

	def getDataDecomp(self, sx = 'q'):
		'''
		obtain the three terms of the decompostion
		'''

		temp1 = self.dxData[sx][['r2dxdn', 'r1dxdn_emp',
			'r1dxdp_q_emp', 'u1_dxdy']]
		temp2 = self.pData[['pi_nq', 'sp_n_p_n']]
		temp3 = self.dnData[['u1_dndy']]
		temp4 = self.xData[[sx]]

		dData = pd.concat([temp1, temp2, temp3, temp4], axis = 1)

		dData['n'] = self.FR2.n

		dData['total'] = dData['r2dxdn']
		dData['price'] = dData['r1dxdp_q_emp']*dData['pi_nq']
		dData['substitution'] = dData['r1dxdn_emp'] \
			*(1-dData['sp_n_p_n'] * dData['u1_dndy'])
		dData['income'] = dData['u1_dxdy']*dData['sp_n_p_n']

		self.dData = dData

		return dData

	def getDataDecomp_e(self, sx = 'q'):
		'''
		obtain the three terms of the decompostion in quasi-elasticity form
		'''

		temp1 = self.dxData[sx][['r2dxdn', 'r1dxdn_emp',
			'r1dxdp_q_emp', 'u1_dxdy']]
		temp2 = self.pData[['pi_nq', 'sp_n_p_n']]
		temp3 = self.dnData[['u1_dndy']]
		temp4 = self.xData[[sx]]

		dData_e = pd.concat([temp1, temp2, temp3, temp4], axis = 1)

		dData_e['n'] = self.FR2.n

		dData_e['total'] = dData_e['r2dxdn'] / dData_e[sx]
		dData_e['price'] = dData_e['r1dxdp_q_emp']*dData_e['pi_nq'] / dData_e[sx]
		dData_e['substitution'] = dData_e['r1dxdn_emp'] \
			*(1-dData_e['sp_n_p_n'] * dData_e['u1_dndy']) / dData_e[sx]
		dData_e['income'] = dData_e['u1_dxdy']*dData_e['sp_n_p_n'] / dData_e[sx]

		self.dData_e = dData_e

		return dData_e









